import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
import json
import cv2
import mediapipe as mp
import requests

class PoseEditor:
    def __init__(self, landmarks, image_shape):
        self.landmarks = landmarks
        self.image_shape = image_shape
        self.fig, self.ax = plt.subplots()
        self.circles = []
        self.lines = []
        self.selected_point = None
        self.draw_pose()

    def draw_pose(self):
        """
        绘制骨架图。
        """
        self.ax.clear()
        self.ax.set_xlim(0, 1)
        self.ax.set_ylim(1, 0)  # 翻转 Y 轴
        self.ax.set_aspect("equal")

        # 绘制点和骨架
        for i, lm in enumerate(self.landmarks):
            circle = plt.Circle((lm["x"], lm["y"]), 0.01, color="blue", picker=True)
            self.ax.add_artist(circle)
            self.circles.append((i, circle))

        # 添加连接线（示例：从肩膀到肘部）

        connections = [
            (11, 12), (23, 24), (11, 23), (12, 24),  # 身体主干
            (11, 13), (13, 15), (12, 14), (14, 16),  # 手臂
            (23, 25), (25, 27), (24, 26), (26, 28),  # 腿
            (15, 21), (16, 22), (27, 31), (28, 32),  # 手腕到手指，脚踝到脚尖
            (32, 30),(29, 31),(28, 30),(27, 29)
        ]

        for p1, p2 in connections:
            line = self.ax.plot(
                [self.landmarks[p1]["x"], self.landmarks[p2]["x"]],
                [self.landmarks[p1]["y"], self.landmarks[p2]["y"]],
                "r-"
            )
            self.lines.append(line)

        plt.draw()

    def on_pick(self, event):
        """
        鼠标选中点时的回调。
        """
        artist = event.artist
        for i, circle in self.circles:
            if artist == circle:
                self.selected_point = i
                break

    def on_motion(self, event):
        """
        鼠标拖动点时的回调。
        """
        if self.selected_point is not None:
            self.landmarks[self.selected_point]["x"] = event.xdata
            self.landmarks[self.selected_point]["y"] = event.ydata
            self.draw_pose()

    def on_release(self, event):
        """
        鼠标释放时的回调。
        """
        self.selected_point = None

    def save_landmarks(self, output_path):
        """
        保存修改后的关键点。
        """
        with open(output_path, "w") as f:
            json.dump(self.landmarks, f, indent=4)


def get_pose_data(image_path):
    """
    使用 MediaPipe 提取姿势数据，同时返回图片尺寸。
    """
    mp_pose = mp.solutions.pose
    pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5)

    # 读取图片
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Image not found at {image_path}")
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 获取姿势关键点
    results = pose.process(image_rgb)
    if not results.pose_landmarks:
        raise ValueError("No pose landmarks detected!")

    # 提取关键点坐标
    landmarks = [
        {"x": lm.x, "y": lm.y, "z": lm.z, "visibility": lm.visibility}
        for lm in results.pose_landmarks.landmark
    ]
    return landmarks, image.shape

"""
把远程图片下载到本地
"""
def download_image(url, save_path):
    response = requests.get(url, stream=True)
    if response.status_code == 200:
        with open(save_path, 'wb') as f:
            for chunk in response.iter_content(1024):
                f.write(chunk)
        print(f"Image downloaded to {save_path}")
        return save_path
    else:
        print(f"Failed to download image: {response.status_code}")
        return None


image_url = "http://tj.yundongchufang.net/statics/upload/admup/4952.png"  # 替换为您的远程图片 URL
image_path = "./ziyuan/zishiys.png"  # 替换为您的图片路径
download_image(image_url, image_path)

landmarks, shape = get_pose_data(image_path)  # 获取关键点和图片尺寸

original_image=cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
# 转换为 RGB 格式（Matplotlib 支持）
if original_image.shape[2] == 4:  # RGBA
    rgba_image = cv2.cvtColor(original_image, cv2.COLOR_BGRA2RGBA)
else:
    rgba_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)


plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号 '-' 显示为乱码的问题
plt.figure(figsize=(5, 5))
plt.subplot(1, 2, 1)
plt.title("远程服务器上面姿势图片")
plt.imshow(rgba_image)
plt.axis("off")



# 启动编辑器
editor = PoseEditor(landmarks, shape)
editor.fig.canvas.mpl_connect("pick_event", editor.on_pick)
editor.fig.canvas.mpl_connect("motion_notify_event", editor.on_motion)
editor.fig.canvas.mpl_connect("button_release_event", editor.on_release)
plt.show()

# 保存修改结果
editor.save_landmarks("edited_pose_data.json")
print("Edited pose data saved.")
